# MCMC types and data structures for FLEXINVERT Bayesian inversion
# Inspired by MCIntegration.jl MCMC framework and FLEXINVERT documentation

module MCMCTypes

using LinearAlgebra
using Random
using Statistics

export MCMCState, MCMCConfiguration, MCMCOptions, MCMCResult, ProposalType
export ChainDiagnostics, HyperParameters, PosteriorSample
export acceptance_rate, update_acceptance!, reset_acceptance!

"""
    ProposalType

Enumeration of available MCMC proposal types.
"""
@enum ProposalType begin
    MetropolisHastings    # Standard Metropolis-Hastings
    CrankNicolson        # Preconditioned Crank-Nicolson (pCN)
    MALA                 # Metropolis-adjusted Langevin algorithm
    BlockGibbs           # Block-wise Metropolis-within-Gibbs
    HamiltonianMC        # Hamiltonian Monte Carlo (future)
end

"""
    HyperParameters

Hierarchical parameters for observation and prior error modeling.
"""
mutable struct HyperParameters
    # Observation error scaling factors (one per observation class)
    obs_error_scale::Vector{Float64}

    # Prior covariance parameters
    prior_scale::Float64
    correlation_length::Float64

    # CAR model parameters for spatial residuals
    car_precision::Float64
    car_correlation::Float64

    # Student-t degrees of freedom (if using robust likelihood)
    nu_obs::Float64

    HyperParameters() = new(
        Float64[],      # obs_error_scale
        1.0,            # prior_scale
        100.0,          # correlation_length (km)
        1.0,            # car_precision
        0.5,            # car_correlation
        4.0             # nu_obs
    )
end

"""
    MCMCState{T}

Current state of MCMC chain tracking selected variables and acceptance statistics.
Similar to _State in MCIntegration.jl.
"""
mutable struct MCMCState{T<:AbstractFloat}
    # Current state vectors
    x_chi::Vector{T}                    # State in decorrelated (chi) space
    x_phys::Vector{T}                   # State in physical space

    # Current evaluation
    log_posterior::T                    # Current log-posterior value
    log_likelihood::T                   # Current log-likelihood
    log_prior::T                        # Current log-prior
    gradient::Vector{T}                 # Current gradient (if available)

    # Proposal statistics
    n_proposed::Int                     # Number of proposals made
    n_accepted::Int                     # Number of proposals accepted

    # Hyperparameters
    hyperparams::HyperParameters        # Current hyperparameter values

    # Block update info
    current_block::Int                  # Current block being updated (0 = full state)

    function MCMCState{T}(n_vars::Int) where T
        new{T}(
            zeros(T, n_vars),           # x_chi
            zeros(T, n_vars),           # x_phys
            -Inf,                       # log_posterior
            -Inf,                       # log_likelihood
            -Inf,                       # log_prior
            zeros(T, n_vars),           # gradient
            0,                          # n_proposed
            0,                          # n_accepted
            HyperParameters(),          # hyperparams
            0                           # current_block
        )
    end
end

MCMCState(n_vars::Int) = MCMCState{Float64}(n_vars)

"""
    MCMCConfiguration

Configuration for MCMC sampling including proposal settings and adaptation.
"""
struct MCMCConfiguration
    # Chain settings
    n_samples::Int                      # Total number of samples
    n_burnin::Int                       # Burn-in samples
    n_thin::Int                         # Thinning interval

    # Proposal settings
    proposal_type::ProposalType         # Primary proposal mechanism
    step_size::Float64                  # Base step size / scaling
    block_structure::Vector{Vector{Int}} # Block definitions for grouped updates

    # Adaptation settings
    adapt_step_size::Bool               # Adapt step size during burn-in
    target_acceptance::Float64          # Target acceptance rate
    adaptation_window::Int              # Window for computing acceptance rates

    # pCN specific settings
    pcn_beta::Float64                   # pCN step size parameter (0 < β < 1)

    # MALA specific settings
    mala_step_size::Float64             # MALA step size

    # Hierarchical sampling
    sample_hyperparams::Bool            # Whether to sample hyperparameters
    hyper_update_freq::Int              # Frequency of hyperparameter updates

    # Diagnostics and output
    compute_diagnostics::Bool           # Compute convergence diagnostics
    save_samples::Bool                  # Save individual samples
    diagnostic_freq::Int                # Frequency of diagnostic computation

    # Parallel settings
    n_chains::Int                       # Number of parallel chains

    function MCMCConfiguration(;
        n_samples::Int = 10000,
        n_burnin::Int = 2000,
        n_thin::Int = 1,
        proposal_type::ProposalType = CrankNicolson,
        step_size::Float64 = 0.1,
        block_structure::Vector{Vector{Int}} = Vector{Int}[],
        adapt_step_size::Bool = true,
        target_acceptance::Float64 = 0.234,  # Optimal for high dimensions
        adaptation_window::Int = 50,
        pcn_beta::Float64 = 0.1,
        mala_step_size::Float64 = 0.01,
        sample_hyperparams::Bool = false,
        hyper_update_freq::Int = 10,
        compute_diagnostics::Bool = true,
        save_samples::Bool = true,
        diagnostic_freq::Int = 100,
        n_chains::Int = 4
    )
        new(n_samples, n_burnin, n_thin, proposal_type, step_size,
            block_structure, adapt_step_size, target_acceptance, adaptation_window,
            pcn_beta, mala_step_size, sample_hyperparams, hyper_update_freq,
            compute_diagnostics, save_samples, diagnostic_freq, n_chains)
    end
end

"""
    ChainDiagnostics

Convergence diagnostics for MCMC chains.
"""
mutable struct ChainDiagnostics
    # Gelman-Rubin diagnostic
    rhat::Vector{Float64}               # R̂ statistic per parameter

    # Effective sample sizes
    ess::Vector{Float64}                # Effective sample size per parameter

    # Autocorrelation
    autocorr_time::Vector{Float64}      # Autocorrelation time per parameter

    # Acceptance rates
    acceptance_rate::Float64            # Overall acceptance rate
    block_acceptance_rates::Vector{Float64} # Per-block acceptance rates

    # Geweke diagnostics
    geweke_z::Vector{Float64}           # Geweke Z-scores

    ChainDiagnostics() = new(
        Float64[], Float64[], Float64[], 0.0, Float64[], Float64[]
    )
end

"""
    PosteriorSample

Single sample from the posterior distribution.
"""
struct PosteriorSample{T<:AbstractFloat}
    x_phys::Vector{T}                   # State in physical space
    x_chi::Vector{T}                    # State in chi space
    log_posterior::T                    # Log-posterior value
    log_likelihood::T                   # Log-likelihood value
    log_prior::T                        # Log-prior value
    hyperparams::HyperParameters        # Hyperparameter values
    iteration::Int                      # Sample iteration number
end

"""
    MCMCResult

Complete results from MCMC sampling.
"""
struct MCMCResult{T<:AbstractFloat}
    # Samples
    samples::Vector{PosteriorSample{T}} # Posterior samples (after burn-in and thinning)

    # Chain states
    final_states::Vector{MCMCState{T}}  # Final state of each chain

    # Summary statistics
    posterior_mean::Vector{T}           # Posterior mean
    posterior_std::Vector{T}            # Posterior standard deviation
    credible_intervals::Matrix{T}       # 95% credible intervals (2 x n_params)

    # Diagnostics
    diagnostics::ChainDiagnostics       # Convergence diagnostics

    # Metadata
    n_chains::Int                       # Number of chains run
    n_samples_total::Int                # Total samples per chain
    n_samples_kept::Int                 # Samples kept after burn-in/thinning
    sampling_time::Float64              # Total sampling time (seconds)

    # Configuration used
    config::MCMCConfiguration          # Configuration used for sampling
end

"""
    acceptance_rate(state::MCMCState)

Compute current acceptance rate from MCMC state.
"""
function acceptance_rate(state::MCMCState)
    state.n_proposed == 0 ? 0.0 : state.n_accepted / state.n_proposed
end

"""
    update_acceptance!(state::MCMCState, accepted::Bool)

Update acceptance statistics in MCMC state.
"""
function update_acceptance!(state::MCMCState, accepted::Bool)
    state.n_proposed += 1
    if accepted
        state.n_accepted += 1
    end
end

"""
    reset_acceptance!(state::MCMCState)

Reset acceptance statistics (useful for adaptation windows).
"""
function reset_acceptance!(state::MCMCState)
    state.n_proposed = 0
    state.n_accepted = 0
end

end # module MCMCTypes
